package com.jim.healio.graphrag.service;

import com.jim.healio.graphrag.model.MedicalEntity;
import com.jim.healio.graphrag.model.MedicalRelationship;
import org.neo4j.driver.Driver;
import org.neo4j.driver.Session;
import org.neo4j.driver.Transaction;
import org.neo4j.driver.Record;
import org.neo4j.driver.Result;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.ArrayList;
import java.util.Map;
import java.util.HashMap;

import static org.neo4j.driver.Values.parameters;

/**
 * Neo4j图数据库服务
 * 用于存储和查询医学知识图谱
 */
@Service
public class Neo4jGraphService {
    
    private static final Logger logger = LoggerFactory.getLogger(Neo4jGraphService.class);
    
    @Autowired(required = false)
    private Driver neo4jDriver;
    
    /**
     * 存储医学实体
     */
    public void storeEntity(MedicalEntity entity) {
        if (neo4jDriver == null) {
            logger.warn("Neo4j驱动未配置，跳过实体存储: {}", entity.getName());
            return;
        }
        
        try (Session session = neo4jDriver.session()) {
            session.executeWrite(tx -> {
                String cypher = """
                    MERGE (e:MedicalEntity {id: $id})
                    SET e.name = $name,
                        e.type = $type,
                        e.description = $description,
                        e.aliases = $aliases,
                        e.confidence = $confidence,
                        e.createdAt = $createdAt,
                        e.properties = $properties
                    RETURN e
                    """;
                
                Map<String, Object> parameters = new HashMap<>();
                parameters.put("id", entity.getId());
                parameters.put("name", entity.getName());
                parameters.put("type", entity.getType().name());
                parameters.put("description", entity.getDescription());
                parameters.put("aliases", entity.getAliases());
                parameters.put("confidence", entity.getConfidence());
                parameters.put("createdAt", entity.getCreatedAt());
                parameters.put("properties", entity.getProperties());
                
                tx.run(cypher, parameters);
                logger.info("实体存储成功: {}", entity.getName());
                return null;
            });
        } catch (Exception e) {
            logger.error("实体存储失败: " + entity.getName(), e);
            throw new RuntimeException("实体存储失败", e);
        }
    }
    
    /**
     * 存储医学关系
     */
    public void storeRelationship(MedicalRelationship relationship) {
        if (neo4jDriver == null) {
            logger.warn("Neo4j驱动未配置，跳过关系存储");
            return;
        }
        
        try (Session session = neo4jDriver.session()) {
            session.executeWrite(tx -> {
                String cypher = """
                    MATCH (source:MedicalEntity {id: $sourceId})
                    MATCH (target:MedicalEntity {id: $targetId})
                    MERGE (source)-[r:%s {id: $relationshipId}]->(target)
                    SET r.description = $description,
                        r.strength = $strength,
                        r.createdAt = $createdAt,
                        r.source = $source,
                        r.properties = $properties
                    RETURN r
                    """.formatted(relationship.getType().name());
                
                Map<String, Object> parameters = new HashMap<>();
                parameters.put("sourceId", relationship.getSourceEntityId());
                parameters.put("targetId", relationship.getTargetEntityId());
                parameters.put("relationshipId", relationship.getId());
                parameters.put("description", relationship.getDescription());
                parameters.put("strength", relationship.getStrength());
                parameters.put("createdAt", relationship.getCreatedAt());
                parameters.put("source", relationship.getSource());
                parameters.put("properties", relationship.getProperties());
                
                tx.run(cypher, parameters);
                logger.info("关系存储成功: {} -> {}", 
                    relationship.getSourceEntityId(), relationship.getTargetEntityId());
                return null;
            });
        } catch (Exception e) {
            logger.error("关系存储失败", e);
            throw new RuntimeException("关系存储失败", e);
        }
    }
    
    /**
     * 查询医学实体
     */
    public List<MedicalEntity> queryEntities(String query) {
        if (neo4jDriver == null) {
            logger.warn("Neo4j驱动未配置，返回空结果");
            return new ArrayList<>();
        }
        
        try (Session session = neo4jDriver.session()) {
            return session.executeRead(tx -> {
                String cypher = """
                    MATCH (e:MedicalEntity)
                    WHERE e.name CONTAINS $query 
                       OR e.description CONTAINS $query
                       OR any(alias IN e.aliases WHERE alias CONTAINS $query)
                    RETURN e
                    ORDER BY e.confidence DESC
                    LIMIT 10
                    """;
                
                Result result = tx.run(cypher, parameters("query", query));
                List<MedicalEntity> entities = new ArrayList<>();
                
                while (result.hasNext()) {
                    Record record = result.next();
                    MedicalEntity entity = convertToMedicalEntity(record.get("e").asNode());
                    entities.add(entity);
                }
                
                logger.info("查询到 {} 个实体", entities.size());
                return entities;
            });
        } catch (Exception e) {
            logger.error("实体查询失败: " + query, e);
            return new ArrayList<>();
        }
    }
    
    /**
     * 分析症状关联
     */
    public List<MedicalEntity> analyzeSymptomRelations(List<String> symptoms) {
        if (neo4jDriver == null) {
            logger.warn("Neo4j驱动未配置，返回空结果");
            return new ArrayList<>();
        }
        
        try (Session session = neo4jDriver.session()) {
            return session.executeRead(tx -> {
                String cypher = """
                    MATCH (symptom:MedicalEntity)-[r:INDICATES|ASSOCIATED_WITH]->(disease:MedicalEntity)
                    WHERE symptom.type = 'SYMPTOM' 
                      AND disease.type = 'DISEASE'
                      AND any(s IN $symptoms WHERE symptom.name CONTAINS s OR any(alias IN symptom.aliases WHERE alias CONTAINS s))
                    RETURN disease, count(r) as relationCount
                    ORDER BY relationCount DESC
                    LIMIT 5
                    """;
                
                Result result = tx.run(cypher, parameters("symptoms", symptoms));
                List<MedicalEntity> diseases = new ArrayList<>();
                
                while (result.hasNext()) {
                    Record record = result.next();
                    MedicalEntity disease = convertToMedicalEntity(record.get("disease").asNode());
                    diseases.add(disease);
                }
                
                logger.info("症状关联分析找到 {} 个可能疾病", diseases.size());
                return diseases;
            });
        } catch (Exception e) {
            logger.error("症状关联分析失败", e);
            return new ArrayList<>();
        }
    }
    
    /**
     * 查询药物相互作用
     */
    public List<MedicalEntity> queryDrugInteractions(List<String> drugs) {
        if (neo4jDriver == null) {
            logger.warn("Neo4j驱动未配置，返回空结果");
            return new ArrayList<>();
        }
        
        try (Session session = neo4jDriver.session()) {
            return session.executeRead(tx -> {
                String cypher = """
                    MATCH (drug1:MedicalEntity)-[r:INTERACTS_WITH]-(drug2:MedicalEntity)
                    WHERE drug1.type = 'DRUG' AND drug2.type = 'DRUG'
                      AND (any(d IN $drugs WHERE drug1.name CONTAINS d OR any(alias IN drug1.aliases WHERE alias CONTAINS d))
                       OR any(d IN $drugs WHERE drug2.name CONTAINS d OR any(alias IN drug2.aliases WHERE alias CONTAINS d)))
                    RETURN drug1, drug2, r
                    ORDER BY r.strength DESC
                    """;
                
                Result result = tx.run(cypher, parameters("drugs", drugs));
                List<MedicalEntity> interactions = new ArrayList<>();
                
                while (result.hasNext()) {
                    Record record = result.next();
                    
                    // 创建表示相互作用的虚拟实体
                    String drug1Name = record.get("drug1").asNode().get("name").asString();
                    String drug2Name = record.get("drug2").asNode().get("name").asString();
                    String description = record.get("r").asRelationship().get("description").asString("");
                    
                    MedicalEntity interaction = new MedicalEntity(
                        "interaction_" + drug1Name + "_" + drug2Name,
                        drug1Name + " - " + drug2Name + " 相互作用",
                        MedicalEntity.EntityType.DRUG,
                        description.isEmpty() ? "存在药物相互作用" : description
                    );
                    
                    // 添加风险等级信息
                    double strength = record.get("r").asRelationship().get("strength").asDouble(1.0);
                    if (strength > 0.8) {
                        interaction.getProperties().put("risk_level", "高风险");
                    } else if (strength > 0.5) {
                        interaction.getProperties().put("risk_level", "中风险");
                    } else {
                        interaction.getProperties().put("risk_level", "低风险");
                    }
                    interaction.getProperties().put("monitoring_required", strength > 0.6);
                    
                    interactions.add(interaction);
                }
                
                logger.info("药物相互作用查询找到 {} 个相互作用", interactions.size());
                return interactions;
            });
        } catch (Exception e) {
            logger.error("药物相互作用查询失败", e);
            return new ArrayList<>();
        }
    }
    
    /**
     * 根据路径查找相关实体
     */
    public List<List<MedicalEntity>> findPaths(String startEntity, int maxDepth) {
        if (neo4jDriver == null) {
            logger.warn("Neo4j驱动未配置，返回空结果");
            return new ArrayList<>();
        }
        
        try (Session session = neo4jDriver.session()) {
            return session.executeRead(tx -> {
                String cypher = """
                    MATCH path = (start:MedicalEntity)-[*1..%d]-(end:MedicalEntity)
                    WHERE start.name CONTAINS $query
                    RETURN path
                    LIMIT 20
                    """.formatted(maxDepth);
                
                Result result = tx.run(cypher, parameters("query", startEntity));
                List<List<MedicalEntity>> paths = new ArrayList<>();
                
                while (result.hasNext()) {
                    Record record = result.next();
                    // 这里需要解析路径，暂时简化处理
                    List<MedicalEntity> path = new ArrayList<>();
                    paths.add(path);
                }
                
                logger.info("路径查找完成，找到 {} 条路径", paths.size());
                return paths;
            });
        } catch (Exception e) {
            logger.error("路径查找失败", e);
            return new ArrayList<>();
        }
    }
    
    /**
     * 添加实体-文档关联
     */
    public void addEntityDocumentLink(String entityId, String documentId) {
        if (neo4jDriver == null) {
            logger.warn("Neo4j驱动未配置，跳过实体-文档关联");
            return;
        }
        
        try (Session session = neo4jDriver.session()) {
            session.executeWrite(tx -> {
                String cypher = """
                    MATCH (e:MedicalEntity {id: $entityId})
                    MERGE (d:Document {id: $documentId})
                    MERGE (e)-[r:REFERENCED_IN]->(d)
                    SET r.createdAt = timestamp()
                    RETURN r
                    """;
                
                tx.run(cypher, parameters("entityId", entityId, "documentId", documentId));
                return null;
            });
        } catch (Exception e) {
            logger.error("实体-文档关联失败", e);
        }
    }
    
    /**
     * 创建基础医学关系
     */
    public void createBasicMedicalRelationships() {
        if (neo4jDriver == null) {
            logger.warn("Neo4j驱动未配置，跳过基础关系创建");
            return;
        }
        
        logger.info("创建基础医学关系...");
        
        // 糖尿病相关关系
        createRelationshipIfEntitiesExist("diabetes_type2", "metformin", 
            MedicalRelationship.RelationshipType.TREATS, "二甲双胍是2型糖尿病的一线治疗药物");
        
        createRelationshipIfEntitiesExist("diabetes_type2", "polyuria", 
            MedicalRelationship.RelationshipType.CAUSES, "糖尿病导致多尿症状");
        
        createRelationshipIfEntitiesExist("diabetes_type2", "polydipsia", 
            MedicalRelationship.RelationshipType.CAUSES, "糖尿病导致多饮症状");
        
        // 高血压相关关系
        createRelationshipIfEntitiesExist("hypertension_primary", "acei_drugs", 
            MedicalRelationship.RelationshipType.TREATS, "ACEI类药物用于治疗高血压");
        
        logger.info("基础医学关系创建完成");
    }
    
    /**
     * 如果实体存在则创建关系
     */
    private void createRelationshipIfEntitiesExist(String sourceId, String targetId, 
                                                  MedicalRelationship.RelationshipType type, 
                                                  String description) {
        MedicalRelationship relationship = new MedicalRelationship(sourceId, targetId, type, description);
        try {
            storeRelationship(relationship);
        } catch (Exception e) {
            logger.warn("关系创建失败，可能实体不存在: {} -> {}", sourceId, targetId);
        }
    }
    
    /**
     * 将Neo4j节点转换为MedicalEntity
     */
    private MedicalEntity convertToMedicalEntity(org.neo4j.driver.types.Node node) {
        MedicalEntity entity = new MedicalEntity();
        entity.setId(node.get("id").asString());
        entity.setName(node.get("name").asString());
        entity.setType(MedicalEntity.EntityType.valueOf(node.get("type").asString()));
        entity.setDescription(node.get("description").asString(""));
        entity.setConfidence(node.get("confidence").asDouble(1.0));
        entity.setCreatedAt(node.get("createdAt").asLong(System.currentTimeMillis()));
        
        // 处理别名
        if (node.containsKey("aliases")) {
            entity.setAliases(node.get("aliases").asList(org.neo4j.driver.Value::asString));
        }
        
        // 处理属性
        if (node.containsKey("properties")) {
            entity.setProperties(node.get("properties").asMap());
        }
        
        return entity;
    }
    
    /**
     * 检查Neo4j是否可用
     */
    public boolean isNeo4jAvailable() {
        return neo4jDriver != null;
    }
    
    /**
     * 初始化数据库约束和索引
     */
    public void initializeDatabase() {
        if (neo4jDriver == null) {
            logger.warn("Neo4j驱动未配置，跳过数据库初始化");
            return;
        }
        
        try (Session session = neo4jDriver.session()) {
            session.executeWrite(tx -> {
                // 创建唯一性约束
                try {
                    tx.run("CREATE CONSTRAINT entity_id_unique IF NOT EXISTS FOR (e:MedicalEntity) REQUIRE e.id IS UNIQUE");
                    tx.run("CREATE CONSTRAINT document_id_unique IF NOT EXISTS FOR (d:Document) REQUIRE d.id IS UNIQUE");
                    
                    // 创建索引
                    tx.run("CREATE INDEX entity_name_index IF NOT EXISTS FOR (e:MedicalEntity) ON (e.name)");
                    tx.run("CREATE INDEX entity_type_index IF NOT EXISTS FOR (e:MedicalEntity) ON (e.type)");
                    
                    logger.info("Neo4j数据库初始化完成");
                } catch (Exception e) {
                    logger.warn("数据库初始化时出现警告（可能约束已存在）: {}", e.getMessage());
                }
                return null;
            });
        } catch (Exception e) {
            logger.error("数据库初始化失败", e);
        }
    }
} 